# -*- coding: utf-8 -*-
import xml.parsers.expat
import base64
import os
import tempfile
#from pyPdf import PdfFileWriter, PdfFileReader
from cStringIO import StringIO
from reportlab.pdfgen import canvas
from os import path
from subprocess import Popen, PIPE, STDOUT

"""
    Classe responsável por inserir um selo de autenticidade nos arquivos PDF.
    Além do selo, é possível também adicionar uma string que pode ser utilizada para identificar
    o arquivo em uma base de dados ou sistema de arquivos.
"""
class Carimbo: 
    
    def __init__(self, caminho_selo='/Users/breno/Ifrn/assinatura-digital/selo/selo.jpg'):
        self.caminho_selo = caminho_selo
    """
        Adicionar o selo cujo caminho está definido na constante CAMINHO_ARQUIVO_SELO no arquivo PDF.
        caminho_arquivo: string contendo a localização do arquivo no sistema de arquivos
        identificador: string opcional contendo um identificador para o aquivo que está sendo carimbado
    """
    def carimbar(self, caminho_arquivo, identificador):
        caminho_arquivo_temporario = caminho_arquivo+'.pdf'
        os.rename(caminho_arquivo, caminho_arquivo_temporario)
        print caminho_arquivo_temporario
        arquivo = PdfFileReader(file(caminho_arquivo_temporario, "rb"))
        output = PdfFileWriter()
        pagina0 = arquivo.getPage(0)
        output.addPage(pagina0)
        

        buffer = StringIO()
        p = canvas.Canvas(buffer)
        p.drawString(475, 740, identificador)
        p.drawImage(self.caminho_selo,520,750,41,35) 
        p.showPage()
        p.save()
        pdf = buffer.getvalue()
        buffer.close()
        f = tempfile.NamedTemporaryFile(delete=False)
        f.write(pdf)
        f.close()
        
        selo = PdfFileReader(file(f.name, "rb"))
        pagina0.mergePage(selo.getPage(0))
        for i in range(1, arquivo.getNumPages()):
            output.addPage(arquivo.getPage(i))
        outputStream = file(caminho_arquivo, "wb")
        output.write(outputStream)
        outputStream.close()
        os.remove(f.name)

CABECALHO_XML = """<?xml version="1.0" encoding="UTF-8"?>
<References>
  <Data>"""

RODAPE_XML = """</Data>
<Signature xmlns="http://www.w3.org/2000/09/xmldsig#">
      <SignedInfo>
          <CanonicalizationMethod Algorithm=
           "http://www.w3.org/TR/2001/REC-xml-c14n-20010315"/>
          <SignatureMethod Algorithm=
           "http://www.w3.org/2000/09/xmldsig#rsa-sha1"/>
          <Reference>
          <Transforms>
              <Transform Algorithm=
              "http://www.w3.org/2000/09/xmldsig#enveloped-signature" />
              <Transform Algorithm=
              "http://www.w3.org/TR/1999/REC-xpath-19991116">
                  <XPath>ancestor::Data</XPath>
              </Transform>
          </Transforms>
          <DigestMethod Algorithm=
          "http://www.w3.org/2000/09/xmldsig#sha1"/>
          <DigestValue></DigestValue>
          </Reference>
      </SignedInfo>
      <SignatureValue />
      <KeyInfo>
    <X509Data >
    </X509Data>
      </KeyInfo>
  </Signature>
</References>"""


"""
    Referência para um arquivo XML Security
"""
class ArquivoXMLSec:
    __parser__ = None
    __conteudo__ = None
    __ler__ = False
    __caminho_arquivo_xml__ = None
    """
    Construtor
        diretório: string com o nome completo do diretório (Ex: '/path/')
        nome_arquivo: string com o nome do arquivo (Ex: arquivo.xml)
    """
    def __init__(self, caminho_arquivo_xml):
        self.__parser__ = xml.parsers.expat.ParserCreate();
        self.__parser__.StartElementHandler = self.__start_element__
        self.__parser__.EndElementHandler = self.__end_element__
        self.__parser__.DefaultHandler = self.__default_handler__
        self.__caminho_arquivo_xml__ = caminho_arquivo_xml
        self.__conteudo__ = []
        self.__ler__ = False
        
    def __start_element__(self, name, attrs):
        if name == 'Data':
            self.__ler__ = True
    def __end_element__(self, name):
        if name == 'Data':
            self.__ler__ = False
    
    def __default_handler__(self, data):
        if self.__ler__:
            self.__conteudo__.append(data)   
    
    def __extrair_conteudo_string__(self, xml):
        self.__parser__.Parse(xml, 1)
        conteudo =  ''.join(self.__conteudo__)
        return conteudo
    
    def __extrair_conteudo_arquivo__(self, caminho_arquivo_xml, modo):
        arquivo_xml = open(caminho_arquivo_xml, modo, 0)
        conteudo = None
        try:
            xml = arquivo_xml.read()
            conteudo = self.__extrair_conteudo_string__(xml)
        except Exception,e:
            raise e
        finally:
            arquivo_xml.close()
        return conteudo
    """
        Escreve o arquivo no disco caso ele não exista ou sebrescreve-o
        caso contrário
            conteúdo: bytes que serão encapsulados no arquivo XML
    """
    def criar(self, conteudo, binario=True):
        try:
            arquivo_xml = open(self.__caminho_arquivo_xml__, 'w')
            arquivo_xml.write(CABECALHO_XML)
            if binario:
                conteudo = base64.b64encode(conteudo)
            arquivo_xml.write(conteudo)
            arquivo_xml.write(RODAPE_XML)
        except Exception,e:
            raise e
        finally:
            arquivo_xml.close()
    """
        Incorpora um arquivo presente no disco em um arquivo XMLSec cujo caminho foi especificado no construtor
        contendo seu conteúdo codificado no formato base64 caso seja binário
            caminho_arquivo_binario: string contendo o caminho completo do arquivo no sistema de arquivos
            binario: flag indicando se o arquivo é binário ou não
            remover: flag indicando se o arquivo deve ser removido após a conversão
    """ 
    def incorporar(self, caminho_arquivo_binario, binario=True, remover=True):
        if not path.isfile(caminho_arquivo_binario): raise Exception('O arquivo binário é inválido ou não existe')
        try:
            if binario: modo = 'rb'
            else: modo = 'r'
            arquivo_binario = open(caminho_arquivo_binario, modo)
            bytes = arquivo_binario.read()
            self.criar(bytes, binario)
            if remover:
                os.remove(caminho_arquivo_binario)
        except Exception,e:
            raise e
        finally:
            arquivo_binario.close()      
    """
        Extrai o conteúdo do arquivo previamente inserido em um arquivo XMLSec
            caminho_arquivo_saida: string contendo o caminho completo do arquivo que será gerado
            a partir do conteúdo presente no arquivo XML
    """
    def extrair(self, caminho_arquivo_saida, binario=True):
        if binario:
            modo_escrita = 'wb'
            modo_leitura = 'r'
        else:
            modo_escrita = 'w'
            modo_leitura = 'r'
        try:
            conteudo = self.__extrair_conteudo_arquivo__(self.__caminho_arquivo_xml__, modo_leitura)
            if binario:    
                conteudo = base64.b64decode(conteudo)
            arquivo_binario = open(caminho_arquivo_saida, modo_escrita)
            arquivo_binario.write(conteudo)
        except Exception,e:
            raise e
        finally:
            arquivo_binario.close()  
"""
    Classe responsável por encriptar/decriptar arquivos
"""
class Cifrador:
    """
        Criptografa o arquivo passado como parâmetro com uma determinada senha
            caminho_arquivo: string contendo a localização do arquivo no sistema de arquivos
            senha: string contendo uma senha que será requerida no processo de decriptografia
    """
    def encriptar(self, caminho_arquivo, senha):
        if not path.isfile(caminho_arquivo): raise Exception('O arquivo a ser encriptado é inválido ou não existe')
        cmd = 'openssl enc -aes-256-cbc -salt -in "%s" -out "%s" -pass pass:%s' % (caminho_arquivo, caminho_arquivo+'.tmp', senha)
        out = Popen(cmd, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True).stdout
        retorno = out.read()
        print retorno
        os.rename(caminho_arquivo+'.tmp', caminho_arquivo)
    """
        Decriptografa o arquivo passado como parâmetro com uma determinada senha (deve ser a mesma usada no processo de criptografia)
            caminho_arquivo: string contendo a localização do arquivo no sistema de arquivos
            senha: string contendo a mesma senha utilizada no processo de criptografia
    """
    def decriptar(self, caminho_arquivo, senha):
        if not path.isfile(caminho_arquivo): raise Exception('O arquivo a ser decriptado é inválido ou não existe')
        cmd = 'openssl enc -d -aes-256-cbc -in "%s" -out "%s" -pass pass:%s' % (caminho_arquivo, caminho_arquivo+'.tmp', senha)
        out = Popen(cmd, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True).stdout
        retorno = out.read()
        print retorno
        os.rename(caminho_arquivo+'.tmp', caminho_arquivo)
"""
    Classe responsável por assinar e verificar a autenticidade dos arquivos enviados para o servidor e armazenados
    através do padrão XMLSec.
"""
class CertificadorXML:
    """
        Assina um documento XMLSec utilizando a chave primária da entidade (instituição ou usuário)
            caminho_arq_chave_priv: string contendo o caminho completo do arquivo onde a chave primária está localizada
            caminho_arquivo_xml: string contendo o caminho completo do arquivo XML Sec que será assinado
    """
    def assinar(self, caminho_arq_chave_priv, senha, caminho_arquivo_xml):
        if not path.isfile(caminho_arq_chave_priv): raise Exception('O arquivo contendo a chave privada é inválido ou não existe')
        if not path.isfile(caminho_arquivo_xml): raise Exception('O arquivo XML a ser assinado é inválido ou não existe')
        cmd = 'xmlsec1 sign --privkey "%s" --output "%s" --pwd %s "%s"'%(caminho_arq_chave_priv, caminho_arquivo_xml, senha, caminho_arquivo_xml)
        out = Popen(cmd, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True).stdout
        retorno = out.read()
        if retorno.find('Error:')>=0:
            raise Exception(retorno)
    """
        Verifica a autenticidade de um documento XMLSec utilizando a chave pública da entidade (instituíção ou usuário)
        caminho_arq_chave_pub: string contendo o caminho completo do arquivo onde a chave pública está localizada
        caminho_arquivo_xml: string contendo o caminho completo do arquivo XML Sec que será verificado
        return: True se o arquivo for autêntico, False caso contrário
    """
    def verificar(self, caminho_arq_chave_pub, caminho_arq):
        if not path.isfile(caminho_arq_chave_pub): raise Exception('O arquivo contendo a chave pública é inválido ou não existe')
        if not path.isfile(caminho_arq): raise Exception('O arquivo XML a ser verificado é inválido ou não existe')
        cmd = 'xmlsec1 verify --pubkey "%s" "%s"' % (caminho_arq_chave_pub, caminho_arq)
        out = Popen(cmd, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True).stdout
        retorno = out.read()
        return retorno.find('OK')>=0
    """
        Verifica a autenticidade de um documento XMLSec utilizando o certificado da entidade certificadora
            caminho_arq_cert_entidade: string contendo o caminho do arquivo do certificado da entidade certificadora
            caminho_arquivo_xml: string contendo o caminho do arquivo XMLSec que será verificado
            
            return: True se o arquivo for autêntico, False caso contrário
    """
    def verificar_com_certificado(self, caminho_arq_cert_entidade, caminho_arquivo_xml):
        if not path.isfile(caminho_arq_cert_entidade): raise Exception('O arquivo contendo o certificado é inválido ou não existe')
        if not path.isfile(caminho_arquivo_xml): raise Exception('O arquivo XML a ser verificado é inválido ou não existe')
        cmd1 = 'openssl x509 -inform pem -in "%s" -pubkey -noout' % caminho_arq_cert_entidade
        out1 = Popen(cmd1, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True).stdout
        conteudo = out1.read()
        arq_chave_pub = tempfile.NamedTemporaryFile(delete=False)
        arq_chave_pub.write(conteudo)
        arq_chave_pub.close()
        autentico = self.verificar(arq_chave_pub.name, caminho_arquivo_xml)
        os.unlink(arq_chave_pub.name)
        return autentico

"""
    Classe responsável por calcular a assinatura/hash dos arquivos enviados para o servidor e verificar sua autenticidade.
    Diferente dos arquivos armazenados através do padrão XMLSec, os arquivos assinados/verificados por esta classes podem
    ser armazenados em uma base de dados ou qualquer outro mecanismo que possa armazenar tanto o arquivo binário quanto sua
    assinatura.
"""
class CertificadorBD:
    """
        Calcula o hash de um determinando arquivo presente do disco rígido. Utiliza o algorítmo sha256
            caminho_arquivo: string contendo o caminho do arquivo cujo hash será calculado
    """
    def calcular_hash(self, caminho_arquivo):
        if not path.isfile(caminho_arquivo): raise Exception('O arquivo cujo hash deseja-se calcular é inválido ou não existe')
        cmd1 = 'openssl dgst -sha256 < "%s"' % (caminho_arquivo)
        out1 = Popen(cmd1, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True).stdout
        hash = out1.read()
        if hash.startswith('(stdin)= '): return hash[9:] #compatibilidade com MAC OS
        else: return hash
    """
        Calcula e retorna a assinatura de um arquivo presente no disco rígido baseado na chave privada.
        A senha deve ser fornecida caso a chave privada seja projetida.
            caminho_arquivo_chave_priv: string contendo o caminho do arquivo da chave privada
            senha: senha utilizada no processo de criação da chave. Caso nenhuma senha tenha sido utilizada, o valor desse parâmetro é ignorado
            caminho_arquivo: string contendo o caminho do arquivo cuja assinatura será extraída
    """
    def calcular_assinatura(self, caminho_arquivo_chave_priv, senha, caminho_arquivo):
        if not path.isfile(caminho_arquivo_chave_priv): raise Exception('O arquivo contendo a chave privada é inválido ou não existe')
        if not path.isfile(caminho_arquivo): raise Exception('O arquivo a ser assinado é inválido ou não existe')
        hash = self.calcular_hash(caminho_arquivo)
        arq_hash = tempfile.NamedTemporaryFile(delete=False)
        arq_hash.write(hash)
        arq_hash.close()
        cmd2 = 'openssl rsautl -sign -inkey "%s" -keyform PEM -in "%s" -passin pass:%s' % (caminho_arquivo_chave_priv, arq_hash.name, senha)
        out2 = Popen(cmd2, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True).stdout
        retorno = out2.read()
        if retorno.find('unable')!=-1:
            raise Exception(retorno)
            os.remove(arq_hash.name)
        os.remove(arq_hash.name)
        return base64.b64encode(retorno)
    """
        Verifica a autenticidade e integridade de um arquivo presente no disco rígido baseado na assinatura previamente calculada
        e na chave pública associada à chave privada utilizada no processo de assinatura
            caminho_arquivo_chave_pub: string contendo o caminho do arquivo da chave publica
            assinatura_64: string contendo a assinatura do arquivo codificado no formato base64
            caminho_arquivo: string contendo o caminho do arquivo cuja verificação será realizada
            return: True se o arquivo for autêntico, False caso contrário
    """        
    def verificar(self, caminho_arquivo_chave_pub, assinatura_64, caminho_arquivo):
        if not path.isfile(caminho_arquivo_chave_pub): raise Exception('O arquivo contendo a chave pública é inválido ou não existe')
        if not path.isfile(caminho_arquivo): raise Exception('O arquivo a ser verificado é inválido ou não existe')
        arq_assinatura = tempfile.NamedTemporaryFile(delete=False)
        try:
            arq_assinatura.write(base64.b64decode(assinatura_64))
        except Exception:
            arq_assinatura.close()
            os.remove(arq_assinatura.name)
            return False
        arq_assinatura.close()
        cmd2 = 'openssl rsautl -verify -inkey "%s" -keyform PEM -pubin -in "%s"' % (caminho_arquivo_chave_pub, arq_assinatura.name)
        out2 = Popen(cmd2, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True).stdout
        verificacao = out2.read()
        os.remove(arq_assinatura.name)
        hash = self.calcular_hash(caminho_arquivo)
        return verificacao == hash
    """
        Verifica a autenticidade e integridade de um arquivo presente no disco rígido baseado na assinatura previamente calculada
        e no certificado (formato PEM) associado à chave privada utilizada no processo de assinatura
            caminho_arquivo_certificado: string contendo o caminho do arquivo do certificado (formato PEM)
            assinatura_64: string contendo a assinatura do arquivo codificado no formato base64
            caminho_arquivo: string contendo o caminho do arquivo cuja verificação será realizada
            return: True se o arquivo for autêntico, False caso contrário
    """ 
    def verificar_com_certificado(self, caminho_certificado, assinatura_64, caminho_arquivo):
        if not path.isfile(caminho_certificado): raise Exception('O arquivo de certificado é inválido ou não existe')
        if not path.isfile(caminho_arquivo): raise Exception('O arquivo a ser verificado é inválido ou não existe')
        cmd1 = 'openssl x509 -inform pem -in "%s" -pubkey -noout' % caminho_certificado
        out1 = Popen(cmd1, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True).stdout
        conteudo = out1.read()
        arq_chave_pub = tempfile.NamedTemporaryFile(delete=False)
        arq_chave_pub.write(conteudo)
        arq_chave_pub.close()        
        autentico = self.verificar(arq_chave_pub.name, assinatura_64, caminho_arquivo)       
        os.remove(arq_chave_pub.name)
        return autentico
